#!/bin/bash

# 脚本名称：list_disk_info.sh

# 功能说明：
# 本脚本用于列出虚拟磁盘或物理磁盘的磁盘信息。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 显示目标设备的详细信息，包括磁盘大小、分区表类型和分区信息。
# 5. 使用 trap 机制在脚本退出时自动断开挂载的 nbd 设备，确保系统资源得到及时释放。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted。
# - 在执行磁盘操作前，建议备份重要数据，以防万一。

# 安装依赖包
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe")
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted..."
    sudo apt-get install -y -qq qemu-utils parted >/dev/null
}

# 全局变量
disk=""
nbd_device=""
target_device=""
virtual_disk_attached=false

# 清理函数：如果虚拟磁盘已挂载，则断开 nbd 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "清理：断开 nbd 设备 $nbd_device"
        sudo qemu-nbd -d "$nbd_device"
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 查找空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘
mount_virtual_disk() {
    sudo modprobe nbd max_part=8
    nbd_device=$(select_free_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi
    echo "将虚拟磁盘镜像挂载到 nbd 设备：$nbd_device"
    sudo qemu-nbd -c "$nbd_device" "$disk"
    sudo udevadm settle
    virtual_disk_attached=true
}

# 获取用户输入
get_disk_input() {
    if [ -n "$1" ]; then
        disk="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " disk
    fi
}

# 处理设备类型
process_disk() {
    if [ -b "$disk" ]; then
        echo "选择的物理磁盘: $disk"
        target_device="$disk"
    elif [ -f "$disk" ]; then
        echo "选择的虚拟磁盘镜像文件: $disk"
        mount_virtual_disk
        target_device="$nbd_device"
    else
        echo "错误： '$disk' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 显示设备信息
display_disk_info() {
    echo "--------------------------------"
    echo "磁盘: $target_device"
    echo "--------------------------------"

    # 获取磁盘大小
    local disk_size
    disk_size=$(sudo lsblk -dn -o SIZE "$target_device" 2>/dev/null)
    echo "磁盘大小: $disk_size"

    # 获取分区表类型
    local part_table
    part_table=$(sudo parted -s "$target_device" print | grep "Partition Table:" | awk '{print $NF}')
    echo "分区表类型: $part_table"

    # 列出分区信息
    echo "分区信息:"
    sudo lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG "$target_device"

    echo "--------------------------------"
}

# 主函数
main() {
    install_dependencies
    get_disk_input "$1"
    process_disk
    display_disk_info
}

main "$@"
